#include <iostream>
#include <openvdb/openvdb.h>
#include <openvdb/tools/LevelSetSphere.h>
#include <openvdb/points/PointCount.h>

#include "include/timer.h"

struct ReadValueOp
{
    explicit ReadValueOp(openvdb::Index64 index) : mIndex(index) { }
    void operator()(const openvdb::tree::LeafManager<
        openvdb::points::PointDataTree>::LeafRange& range) const {
        for (auto leafIter = range.begin(); leafIter; ++leafIter) {
            for (auto indexIter = leafIter->beginIndexOn();
                indexIter; ++indexIter) {
                const openvdb::points::AttributeArray& array =
                    leafIter->constAttributeArray(mIndex);
                openvdb::points::AttributeHandle<openvdb::Vec3f> handle(array);
                openvdb::Vec3f value = handle.get(*indexIter);
            }
        }
    }
    openvdb::Index64 mIndex;
};


int main()
{
    // Initialize grid types and point attributes types.
    openvdb::initialize();
    // Generate a level set grid.
    openvdb::FloatGrid::Ptr sphereGrid = openvdb::tools::createLevelSetSphere<openvdb::FloatGrid>(20.0, openvdb::Vec3f(1.5, 2, 3), 0.5);
    // Use the topology to create a PointDataTree
    openvdb::points::PointDataTree::Ptr pointTree(new openvdb::points::PointDataTree(sphereGrid->tree(), 0, openvdb::TopologyCopy()));

    // Define the position type and codec using fixed-point 16-bit compression.
    using PositionAttribute = openvdb::points::TypedAttributeArray<openvdb::Vec3f, openvdb::points::FixedPointCodec<false>>;
    openvdb::NamePair positionType = PositionAttribute::attributeType();

    // Create a new Attribute Descriptor with position only
    openvdb::points::AttributeSet::Descriptor::Ptr descriptor(openvdb::points::AttributeSet::Descriptor::create(positionType));
    // Determine the number of points / voxel and points / leaf.
    openvdb::Index pointsPerVoxel = 8;
    openvdb::Index voxelsPerLeaf = openvdb::points::PointDataGrid::TreeType::LeafNodeType::SIZE;
    openvdb::Index pointsPerLeaf = pointsPerVoxel * voxelsPerLeaf;

    // Create the points grid.
    openvdb::points::PointDataGrid::Ptr points = openvdb::points::PointDataGrid::create(pointTree);
    // Set the name of the grid.
    points->setName("Points");
    points->setTransform(sphereGrid->transform().copy());
    // Iterate over the leaf nodes in the point tree.
    for (auto leafIter = pointTree->beginLeaf(); leafIter; ++leafIter) {
        // Initialize the attributes using the descriptor and point count.
        leafIter->initializeAttributes(descriptor, pointsPerLeaf);
        // Initialize the voxel offsets
        openvdb::Index offset(0);
        for (openvdb::Index index = 0; index < voxelsPerLeaf; ++index) {
            offset += pointsPerVoxel;
            leafIter->setOffsetOn(index, offset);
        }
    }

    TIC;
    // Read
    for (auto leafIter = pointTree->beginLeaf(); leafIter; ++leafIter) {
        openvdb::points::AttributeArray& array = leafIter->attributeArray("P");
        openvdb::points::AttributeWriteHandle<openvdb::Vec3f> handle(array); // AttributeHandle
        // Iterate over active indices in the leaf.
        for (auto indexIter = leafIter->beginIndexOn(); indexIter; ++indexIter) {
            // Retrieve value
            openvdb::Vec3f value = handle.get(*indexIter);
            // std::cout << indexIter << " : value of point" << value << std::endl;
            // Set value
            handle.set(*indexIter, openvdb::Vec3f(1.5, 2, 3));
            value = handle.get(*indexIter);
            // std::cout << " after set value of point" << value << std::endl;
            break;
        }
    }
    TOC("normal read and write", true);
    // Create a leaf iterator for the PointDataTree.
    auto leafIter = pointTree->cbeginLeaf();
    // Check that the tree has leaf nodes.
    if (!leafIter) {
        std::cerr << "No Leaf Nodes" << std::endl;
    }
    // Retrieve the index from the descriptor.
    auto descriptor2 = leafIter->attributeSet().descriptor();
    openvdb::Index64 index = descriptor2.find("P");
    // Check that the attribute has been found.
    if (index == openvdb::points::AttributeSet::INVALID_POS) {
        std::cerr << "Invalid Attribute" << std::endl;
    }
    // Create a leaf manager for the points tree.
    openvdb::tree::LeafManager<openvdb::points::PointDataTree> leafManager(*pointTree);
    // Create a new operator
    ReadValueOp op(index);
    // Evaluate in parallel
    tbb::parallel_for(leafManager.leafRange(), op);
    TOC("tbb parallel read", true);

    openvdb::Index count = openvdb::points::pointCount(points->tree());
    
    openvdb::Index64 iterationCount(0);
    double averageY(0.0);
    // Iterate over leaf nodes.
    for (auto leafIter = points->tree().beginLeaf(); leafIter; ++leafIter) {
        // Create a read-only position handle.
        const openvdb::points::AttributeArray& positionArray =
            leafIter->constAttributeArray("P");
        openvdb::points::AttributeHandle<openvdb::Vec3f> positionHandle(
            positionArray);
        // Iterate over the point indices in the leaf.
        for (auto indexIter = leafIter->beginIndexOn(); indexIter; ++indexIter) {
            // Extract the world-space position of the point.
            openvdb::Vec3f voxelPosition = positionHandle.get(*indexIter);
            openvdb::Vec3d xyz = indexIter.getCoord().asVec3d();
            openvdb::Vec3f worldPosition =
                points->transform().indexToWorld(voxelPosition + xyz);
            // Increment the sum.
            averageY += worldPosition.y();
            // Track iteration
            iterationCount++;
        }
    }
    averageY /= double(count);
    std::cout << "IterationCount=" << iterationCount << std::endl;
    std::cout << "AveragePositionInY=" << averageY << std::endl;
}